from typing import Dict, Optional
from chroma_manager import ChromaManager
from model_service import ModelService, QueryResult, create_model_service
from datetime import datetime
import pandas as pd

class VisualizationService:
    def __init__(
        self,
        chroma_manager: ChromaManager,
        model_service_type: str = "langflow",
        **model_kwargs
    ):
        self.chroma = chroma_manager
        self.model_service = create_model_service(model_service_type, **model_kwargs)

    def log(self, message: str, title: str = "Info"):
        print(f"{title}: {message}")

    def generate_kpi_params(self, question: str, limit: int = 10) -> Dict:
        """Generate KPI query parameters from user question"""
        try:
            self.log(question, "Querying indicators for question")
            # Get vector search results
            result = self.chroma.query_indicators(question, limit, True)
            self.log(len(result['indicators']), "Query completed, found indicators")
            
            indicators = self.chroma.format_indicators_table(result['indicators'])
            dimensions = self.chroma.get_dimensions_table(question, limit, result['themes'])

            query_result = QueryResult(indicators, dimensions)
            
            # Generate parameters using model
            return self.model_service.generate_query_params(
                question=question,
                query_result=query_result
            )
        except Exception as e:
            self.log(str(e), "Error in generate_kpi_params")
            raise
        
    def generate_plot_code(
        self, 
        question: str,
        query_params: str,
        df_metadata: str
    ) -> str:
        """Generate Plotly visualization code"""
        return self.model_service.generate_plot_code(
            question=question,
            query_params=query_params,
            df_metadata=df_metadata
        )
    
    def get_plotly_figure(self, plotly_code: str, df: pd.DataFrame):
        """Generate Plotly figure from Plotly code"""
        return self.model_service.get_plotly_figure(
            plotly_code=plotly_code, 
            df=df
        )